from typing import Optional, Union, List, Dict, Any, Callable
from pymilvus import (
    utility,
    FieldSchema, CollectionSchema, DataType, Collection
)
from bson import ObjectId
from .class_func import (
    safe_action, load_default_model, normalize_ids,
    extract_query_result, normalize_fields, normalize_index, dtype_map
)
from .linkMilvus import linkMilvus

'''使用手册：创建数据集合的格式：id, embedding......注意：统一小写,没有后缀s
id将作为主键，禁止自定义名称，应使用uuid或者objectId来保证唯一性，修改和删除时需要用到id。
embedding是自动转换的,作为向量字段,一个数据集合中,只会有这一个向量字段,不得有多个向量字段。

插入数据格式：列表字典，字典中的id将作为主键，text会被自动转换成向量embedding字段。其他键不会被转换。

model_embed参数是模型文本向量化函数，默认使用的是

'''


# ======================================================
#  useMilvus —— 集合操作层（安全封装 + 日志输出）
# ======================================================
class useMilvus:
    def __init__(
        self,
        client: linkMilvus,
        name: str,
        model_embed=None,  # 模型的文本向量化函数
        fields=None,  # 字段定义
        metric_type=None,  # 索引类型
        show: bool = True  # 是否输出日志
    ):
        self.client = client
        self.name = name
        self.alias = client.alias
        self.show = show
        self.model_embed = load_default_model(model_embed)
        self.fields = normalize_fields(fields)
        self.index_params = normalize_index(metric_type)
        self.collection = self._init_collection()

    # ======================================================
    # 向量模型
    # ======================================================

    def _embed(self, texts: List[str], batch_size: int = 64) -> List[List[float]]:
        valid = [t.strip() or "空文本" for t in texts]
        embeddings = []

        for i in range(0, len(valid), batch_size):
            batch = valid[i:i + batch_size]
            emb = self.model_embed(batch)

            if hasattr(emb, "tolist"):
                emb = emb.tolist()
            embeddings.extend(emb)

        return embeddings

    # ======================================================
    # 集合初始化
    # ======================================================

    @safe_action("初始化集合")
    def _init_collection(self) -> Collection:
        if not self.client.has_collection(self.name):
            schema = CollectionSchema(
                fields=self.fields,
                description=f"{self.name} 向量集合"
            )
            coll = Collection(name=self.name, schema=schema, using=self.alias)
            coll.create_index(field_name="embedding",
                              index_params=self.index_params)
        else:
            coll = Collection(name=self.name, using=self.alias)

        coll.load()
        return coll

    # ======================================================
    # 数据操作
    # ======================================================
    @safe_action("插入数据")
    def insertMany(self, data: Union[Dict[str, str], List[Dict[str, str]]]):
        """
        向 Milvus 插入多条数据。
        要求每条记录至少包含 'text' 字段，会自动生成 embedding 并补充到记录中。
        兼容类似 MongoDB 的插入风格。
        """
        # -------- 参数标准化 --------
        if isinstance(data, dict):
            data = [data]
        if not all(isinstance(d, dict) for d in data):
            raise TypeError("参数 data 必须是字典或字典列表")

        # -------- 批量生成 embeddings --------
        texts = [item.pop("text") for item in data]
        embeddings = self._embed(texts)
        for item, emb in zip(data, embeddings):
            item["embedding"] = emb

        # -------- 插入到 Milvus --------
        res = self.collection.insert(data)
        self.collection.flush()
        return res

    @safe_action("删除数据")
    def deleteMany(self, ids: Union[str, ObjectId, List[str], List[ObjectId]]):
        ids = normalize_ids(ids)
        expr = f'id in {ids!r}'
        res = self.collection.delete(expr)
        self.collection.flush()
        return res

    def clear(self):
        # 清空集合
        res = self.collection.delete("id != ''")
        self.collection.flush()
        self.collection.load()
        print(f"[Milvus@{self.alias}/{self.name}]当前集合 {self.name} 已清空")
        return res

    def drop_db(self):
        # 删除集合
        utility.drop_collection(self.name, using=self.alias)
        print(f"[Milvus@{self.alias}/{self.name}]当前集合 {self.name} 已删除")
        return True

    # ======================================================
    # 搜索功能
    # ======================================================
    @safe_action("搜索")
    def find(
        self,
        find_data: str,
        limit: int = 3,
        return_fields: Optional[List[str]] = None,
        min_score: float = 0.0,
    ):
        """
        向量搜索方法
        - find_data: 查询数据（文本或特征）
        - limit: 返回结果数量
        - return_fields: 附加返回字段（自动包含 ['id', 'distance']）
        - min_score: 最小匹配分数（过滤低相似度结果）
        """
        if not find_data.strip():
            return []

        # ---------- 默认输出字段 ----------
        # 永远保证返回 id 和 distance
        base_fields = ["id", "distance"]
        return_fields = list(set(base_fields + (return_fields or [])))
        db_fields = [f for f in return_fields if f not in base_fields]

        # ---------- 文本向量化 ----------
        vector = self._embed([find_data])[0]

        # ---------- 自动探测 metric_type ----------
        try:
            index_info = self.collection.indexes[0].params
            metric_type = index_info.get("metric_type", "COSINE").upper()
        except Exception:
            metric_type = "COSINE"  # fallback 默认文本模式

        # ---------- 自动生成搜索参数 ----------
        search_params = {
            "metric_type": metric_type,
            "params": {},  # AUTOINDEX 通常不需要额外参数
        }

        # ---------- 执行搜索 ----------
        results = self.collection.search(
            data=[vector],
            anns_field="embedding",
            param=search_params,
            limit=limit,
            output_fields=db_fields,
        )

        # ---------- 结果处理 ----------
        hits = []
        for hit in results[0]:
            # 兼容 Milvus 不同版本的字段存取
            fields = getattr(hit, "fields", None) or getattr(
                hit, "entity", {}) or {}
            hit_data = {
                "id": hit.id,
                "distance": hit.distance,
            }
            for f in db_fields:
                if f in fields:
                    hit_data[f] = fields[f]
            hits.append(hit_data)

        # ---------- 过滤 + 排序 ----------
        # COSINE 越大越相似；L2 越小越相似
        reverse_sort = metric_type == "COSINE"

        if reverse_sort:
            filtered = [h for h in hits if h["distance"] >= min_score]
        else:
            # 对于 L2 距离，越小越好，用 max_distance 控制过滤范围
            max_dist = max(h["distance"] for h in hits) if hits else 0.0
            threshold = max_dist * (1 - min_score)
            filtered = [h for h in hits if h["distance"] <= threshold]

        # ---------- 排序并返回 ----------
        return sorted(filtered, key=lambda x: x["distance"], reverse=reverse_sort)

    # ======================================================
    # 更新/插入
    # ======================================================

    @safe_action("插入或更新数据")
    def upsertMany(self, data: Union[Dict[str, str], List[Dict[str, str]]]):
        """
        向 Milvus 插入或更新多条数据。
        如果 ID 已存在，则先删除再插入。
        """
        if isinstance(data, dict):
            data = [data]
        if not all(isinstance(d, dict) for d in data):
            raise TypeError("参数 data 必须是字典或字典列表")

        # 收集所有 ID
        ids = [d.get("id") for d in data if "id" in d]
        if not ids:
            raise ValueError("upsertMany() 需要每条记录包含 'id' 字段")

        # 查询已存在 ID
        existing = self.collection.query(
            expr=f"id in {ids!r}", output_fields=["id"])
        existing_ids = [r["id"] for r in existing] if existing else []

        # 复用 deleteMany() 删除旧数据
        if existing_ids:
            self.deleteMany(existing_ids)

        # 调用 insertMany() 插入新数据
        res = self.insertMany(data)
        self.collection.load()

        return res

    # ======================================================
    # 查询与统计
    # ======================================================
    def count(self) -> int:
        return self.collection.num_entities

    @safe_action("按ID查询", default=[])
    def get(self, ids: Union[str, List[str]]) -> List[Dict[str, Any]]:
        ids = normalize_ids(ids)
        expr = f"id in {ids!r}"
        raw_res = self.collection.query(expr, output_fields=["id"])
        return extract_query_result(raw_res)

    def desc(self):
        """
        打印并返回当前集合的数据结构（schema），包括：
        - 字段名
        - 数据类型（人类可读）
        - 是否主键 / 是否自动ID
        - 维度(dim) 或 最大长度(max_length)
        - 字段描述
        """
        schema = self.collection.schema
        if not schema:
            return []

        field_info = []

        # 收集字段详细信息
        for field in schema.fields:
            params = getattr(field, "params", {}) or {}
            info = {
                "name": field.name,
                "dtype": dtype_map.get(field.dtype, str(field.dtype)),
                "is_primary": field.is_primary,
                "auto_id": field.auto_id,
                "description": field.description or "",
                "dim": params.get("dim"),
                "max_length": params.get("max_length"),
            }
            field_info.append(info)

        # 打印格式化表格
        log_lines = ["\n📘 当前集合结构："]
        for f in field_info:
            dim_info = f"dim={f['dim']}" if f["dim"] else ""
            len_info = f"max_length={f['max_length']}" if f["max_length"] else ""
            extra = ", ".join(x for x in [dim_info, len_info] if x)
            log_lines.append(
                f"  - {f['name']:<15} | {f['dtype']:<13} | 主键={f['is_primary']:<5} | AutoID={f['auto_id']:<5} | {extra} | {f['description']}"
            )
        return field_info

    # ======================================================
    # 说明
    # ======================================================
    def __repr__(self):
        return f'''<useMilvus alias='{self.alias}' name='{self.name}' count={self.count()} host='{self.client.host}:{self.client.port}'>\n'''
